from data.mnist import Mnist
from ch006_adam_optimization.neural_network import NeuralNetwork
import numpy as np
# import matplotlib
# matplotlib.use('qt5agg')
import matplotlib.pyplot as plt

mnist = Mnist()
nn = NeuralNetwork(28 * 28, 100, 30, 10)
costs = nn.train(
    mnist.training_set,
    mnist.dev_set,
    learning_rate=0.003,
    epochs=50,
    mini_batch_size=1024,
    regularization=1.0,
    decay_v=0.9,
    decay_s=0.999)

x = range(0, len(costs))
train_costs = list(map(lambda c: c['train'], costs))
dev_costs = list(map(lambda c: c['dev'], costs))
plt.plot(x, train_costs, 'r-')
plt.plot(x, dev_costs, 'g-')
plt.show()

hit = 0

for (image, label) in mnist.test_set:
    x = image.reshape(28*28, 1)
    predict = nn.predict(x)
    if np.argmax(predict) == np.argwhere(label == 1):
        hit += 1

hit_rate = float(hit) / len(mnist.test_set)

print("Hit rate: {rate}".format(rate=hit_rate))
